{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 粒子群（PSO）算法可视化\n",
    "\n",
    "（参考 https://mp.weixin.qq.com/s/qOkUe0yyfw9Lesn4NU_oPQ）\n",
    "\n",
    "PSO算法最初由James Kennedy和Russell Eberhart于1995年提出。他们假设信息的社会共享将为群体带来优势。当资源分配超出预测时，鱼群中的个体可以从其他成员的发现和经验中受益，而不是争夺食物。关键词是\"社交互动\"。社会行为提高了个体的适应能力，从而使整个群体表现出智慧。个体的适应能力和群体智慧是相互关联的。\n",
    "\n",
    "粒子是搜索空间中的简单实体。我们创建了一个粒子群，并使用问题的目标函数来衡量每个粒子的适应度。然后，粒子根据其个人最佳位置和群体迄今为止的最佳位置，从当前位置移动到下一个位置。通过迭代移动，群体逐渐接近目标函数的最佳点，这个过程会持续几代。\n",
    "\n",
    "从优化算法的角度来看，它有可能克服局域最小值。\n",
    "\n",
    "下面是约定的一些符号：\n",
    "\n",
    "```\n",
    "Number of particles         : i\n",
    "Number of dimensions        : n\n",
    "Fitness function            : f(x_i)\n",
    "Particles                   : x_i = (x_i1, x_i2, ..., x_in)\n",
    "Current velocity            : v_i = (v_i1, v_i2, ..., v_in)\n",
    "Individual particle's best  : p_i = (p_i1, p_i2, ..., p_in)\n",
    "Global particles' best      : p_g = (p_g1, p_g2, ..., p_gn)\n",
    "Inertia component           : w * v_i(t)\n",
    "Cognitive component         : c_1 * r_1 * (p_i - x_i(t))\n",
    "Social component            : c_2 * r_2 * (g_i - x_i(t))\n",
    "Velocity adjustment         : v_i(t+1) <- Inertia+Cognitive+Social\n",
    "Position adjustment         : x_i(t+1) <- x_i(t)+v_i(t+1)\n",
    "```\n",
    "\n",
    "速度调整受 3 个因素影响：先前的速度（惯性分量）、单个粒子的最佳位置（认知分量）和群体的最佳位置（社交分量）。速度是粒子向给定方向移动的速度。粒子的运动受到各个方向上这些权重的影响。系数w称为惯性权重，是使粒子保持与上一代相同方向运动的力。c1和c2是恒定加速度值，其中c1=c2在的原始算法中应用。r1和r2表示超参数，它们会引起一些随机扰动。这些参数值越高，粒子的运动响应越灵敏。我们还假设适应度函数适用于我们案例中的最小化问题。因此，当f(x_i) < f(p_i)时，单个粒子的最佳位置p_i被x_i覆盖。\n",
    "\n",
    "\n",
    "PSO的算法逻辑如下：\n",
    "\n",
    "1. Initialise the particle population array x_i\n",
    "2. Loop\n",
    "3. For each particle, calculate the fitness using the \n",
    "   fitness function f(x_i)\n",
    "4. Compare the current fitness value with its best p_i. \n",
    "   Replace the best with the current value x_i \n",
    "   if it is better than the best.\n",
    "5. Check the swarm’s best particle from individual particle’s best \n",
    "   and assign the best array to the global best p_g.\n",
    "6. Calculate the velocity v_i(t+1) and update the position of \n",
    "   the particles to x_i(t+1)\n",
    "7. If a criterion is met, exit the loop.\n",
    "8. End loop\n",
    "\n",
    "考虑下面的目标函数：\n",
    "\n",
    "$$ f(x, y) = (x-2y+3)^2 + (2x+y-8)^2 $$\n",
    "\n",
    "显然，当两项分别等于 0 时该函数取得极小值 0，反解出 $ (x, y) = (2.6, 2.8) $。\n",
    "下面会利用粒子群算法进行展示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib import animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fitness function\n",
    "# We assume the problem can be expressed by the following equation: \n",
    "# f(x1,x2)=(x1-2*x2+3)^2 + (2*x1+x2-8)^2\n",
    "# The objective is to find a minimum which is 0\n",
    "\n",
    "def fitness_function(x1,x2):\n",
    "    f1=x1-2*x2+3\n",
    "    f2=2*x1+x2-8\n",
    "    z = f1**2+f2**2\n",
    "    return z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_velocity(particle, velocity, pbest, gbest, w_min=0.5, max=1.0, c=0.1):\n",
    "    # Initialise new velocity array\n",
    "    num_particle = len(particle)\n",
    "    new_velocity = np.array([0.0 for _ in range(num_particle)])\n",
    "    \n",
    "    # Randomly generate r1, r2 and inertia weight from normal distribution\n",
    "    r1 = random.uniform(0,max)\n",
    "    r2 = random.uniform(0,max)\n",
    "    w = random.uniform(w_min,max)\n",
    "    c1 = c\n",
    "    c2 = c\n",
    "    # Calculate new velocity\n",
    "    for i in range(num_particle):\n",
    "        new_velocity[i] = w*velocity[i] + c1*r1*(pbest[i]-particle[i])+c2*r2*(gbest[i]-particle[i])\n",
    "    return new_velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_position(particle, velocity):\n",
    "  # Move particles by adding velocity\n",
    "  new_particle = particle + velocity\n",
    "  return new_particle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def pso_2d(population, dimension, position_min, position_max, generation, fitness_criterion):\n",
    "    history = []\n",
    "    \n",
    "    # Initialisation\n",
    "    # Population \n",
    "    particles = [[random.uniform(position_min, position_max) for j in range(dimension)] for i in range(population)]\n",
    "    # Particle's best position\n",
    "    pbest_position = particles\n",
    "    # Fitness\n",
    "    pbest_fitness = [fitness_function(p[0],p[1]) for p in particles]\n",
    "    # Index of the best particle\n",
    "    gbest_index = np.argmin(pbest_fitness)\n",
    "    # Global best particle position\n",
    "    gbest_position = pbest_position[gbest_index]\n",
    "    # Velocity (starting from 0 speed)\n",
    "    velocity = [[0.0 for j in range(dimension)] for i in range(population)]\n",
    "\n",
    "    history = [[list(particles), pbest_fitness,],]\n",
    "\n",
    "    # Loop for the number of generation\n",
    "    for t in range(generation):\n",
    "        # Stop if the average fitness value reached a predefined success criterion\n",
    "        if np.average(pbest_fitness) <= fitness_criterion:\n",
    "            break\n",
    "        else:\n",
    "            for n in range(population):\n",
    "                # Update the velocity of each particle\n",
    "                velocity[n] = update_velocity(particles[n], velocity[n], pbest_position[n], gbest_position)\n",
    "                # Move the particles to new position\n",
    "                particles[n] = update_position(particles[n], velocity[n])\n",
    "            # Calculate the fitness value\n",
    "            pbest_fitness = [fitness_function(p[0],p[1]) for p in particles]\n",
    "            # Find the index of the best particle\n",
    "            gbest_index = np.argmin(pbest_fitness)\n",
    "            # Update the position of the best particle\n",
    "            gbest_position = pbest_position[gbest_index]\n",
    "            \n",
    "            history.append([list(particles), pbest_fitness])\n",
    "\n",
    "    # Print the results\n",
    "    print('Global Best Position: ', gbest_position)\n",
    "    print('Best Fitness Value: ', min(pbest_fitness))\n",
    "    print('Average Particle Best Fitness Value: ', np.average(pbest_fitness))\n",
    "    print('Number of Generation: ', t)\n",
    "    return history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "population = 100\n",
    "dimension = 2\n",
    "position_min = -100.0\n",
    "position_max = 100.0\n",
    "generation = 400\n",
    "fitness_criterion = 10e-4\n",
    "\n",
    "history = pso_2d(population, dimension, position_min, position_max, generation, fitness_criterion)\n",
    "print(len(history))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Plotting prepartion\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "ax.set_zlabel('z')\n",
    "x = np.linspace(position_min, position_max, 80)\n",
    "y = np.linspace(position_min, position_max, 80)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "Z= fitness_function(X,Y)\n",
    "ax.plot_wireframe(X, Y, Z, color='r', linewidth=0.2)\n",
    "\n",
    "# Animation image placeholder\n",
    "images = []\n",
    "\n",
    "# Add plot for each generation (within the generation for-loop)\n",
    "for particles, pbest_fitness in history:\n",
    "    image = ax.scatter3D([p[0] for p in particles],\n",
    "                         [p[1] for p in particles],\n",
    "                         pbest_fitness, c='b')\n",
    "    images.append([image])\n",
    "\n",
    "# Generate the animation image and save\n",
    "animated_image = animation.ArtistAnimation(fig, images)\n",
    "animated_image.save('./pso_simple.gif', writer='pillow')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
